package me.flyray.bsin.server.impl;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.mapping.*;
import co.elastic.clients.elasticsearch.indices.CreateIndexResponse;
import co.elastic.clients.elasticsearch.indices.IndexSettings;
import com.alibaba.cloud.ai.advisor.RetrievalRerankAdvisor;
import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.agent.ReactAgent;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.node.QuestionClassifierNode;
import com.alibaba.cloud.ai.graph.state.AgentStateFactory;
import com.alibaba.cloud.ai.graph.state.strategy.ReplaceStrategy;
import com.alibaba.cloud.ai.memory.redis.RedissonRedisChatMemoryRepository;
import com.alibaba.cloud.ai.model.RerankModel;
import lombok.extern.slf4j.Slf4j;
import me.flyray.bsin.domain.enums.AIAgentChoreographyType;
import me.flyray.bsin.facade.service.AgentChatService;
import me.flyray.bsin.security.contex.LoginInfoContextHelper;
import me.flyray.bsin.security.domain.LoginUser;
import me.flyray.bsin.server.biz.AppAgentEngine;
import me.flyray.bsin.server.biz.RecordingNode;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.shenyu.client.apache.dubbo.annotation.ShenyuDubboService;
import org.apache.shenyu.client.apidocs.annotations.ApiDoc;
import org.apache.shenyu.client.dubbo.common.annotation.ShenyuDubboClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.resolution.ToolCallbackResolver;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.messaging.Message;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static com.alibaba.cloud.ai.graph.action.AsyncEdgeAction.edge_async;
import static com.alibaba.cloud.ai.graph.action.AsyncNodeAction.node_async;
import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

@Slf4j
@ShenyuDubboService(value = "/agentChat/",  timeout = 60000)
public class AgentChatServiceImpl implements AgentChatService {

    private final ChatModel chatModel;
    private final VectorStore vectorStore;
    private final RerankModel rerankModel;
    private final ToolCallbackResolver resolver;
    private final ToolCallbackProvider tools;
    private final MessageWindowChatMemory messageWindowChatMemory;

    @Value("classpath:/prompts/system-qa.st")
    private Resource systemResource;
    private static final String textField = "content";

    @Autowired
    private AppAgentEngine appAgentEngine;

    private final int MAX_MESSAGES = 100;

    public AgentChatServiceImpl(@Qualifier("dashscopeChatModel") ChatModel chatModel, VectorStore vectorStore, RerankModel rerankModel, ElasticsearchClient elasticsearchClient, ToolCallbackResolver resolver, @Qualifier("loadbalancedMcpAsyncToolCallbacks") ToolCallbackProvider tools, RedissonRedisChatMemoryRepository redissonRedisChatMemoryRepository) {
        this.chatModel = chatModel;
        this.vectorStore = vectorStore;
        this.rerankModel = rerankModel;
        this.resolver = resolver;
        this.tools = tools;
        this.messageWindowChatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(redissonRedisChatMemoryRepository)
                .maxMessages(MAX_MESSAGES)
                .build();;

    }

    /**
     * 1、获取编排智能体：智能体类型（编排单智能体、编排多智能体）
     * 2、根据不同类型智能体进行对话结果处理
     * 3、websocket流式处理
     * @param requestMap
     * @return
     */
    @ApiDoc(desc = "chat")
    @ShenyuDubboClient("/chat")
    @Override
    public Map<String, Object> chat(Map<String, Object> requestMap) throws GraphStateException {
        String tenantId = LoginInfoContextHelper.getTenantId();
        String question = MapUtils.getString(requestMap, "question");
        
        // 将tenantId信息拼接到question中，让AI知道当前服务的租户
        String enhancedQuestion = String.format("当前租户ID：%s。用户问题：%s", tenantId, question);
        requestMap.put("tenantId", tenantId);
        requestMap.put("question", enhancedQuestion);

        String aiAgentTypeCode = MapUtils.getString(requestMap, "aiAgentType");
        // 判断类型是否存在和正确
        AIAgentChoreographyType aiAgentType = AIAgentChoreographyType.getInstanceById(aiAgentTypeCode);
        // 判断枚举什么类型，不同的类型做不同的处理
        String answer = null;
        switch (aiAgentType) {
            case SINGLE_AGENT:
                // 单智能体处理逻辑
                answer = appAgentEngine.handleSingleAgent(requestMap);
                break;
            case MULTI_AGENT:
                // 多智能体处理逻辑
                answer =  appAgentEngine.handleMultiAgent(requestMap);
                break;
            default:
                throw new IllegalArgumentException("未知的智能体类型: " + aiAgentType.getDesc());
        }
        Map result = new HashMap<>();
        result.put("answer", answer);

        return result;
    }

    /**
     * reAct chat
     * 1、获取编排智能体：对话智能体、文档智能体、编排智能体、多智能体
     * 2、进行对话结果处理
     * @param requestMap
     * @return
     */
    @ApiDoc(desc = "reActChat")
    @ShenyuDubboClient("/reActChat")
    @Override
    public Map<String, Object> reActChat(Map<String, Object> requestMap) throws GraphStateException, GraphRunnerException {
        String question = MapUtils.getString(requestMap, "question");

        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultToolCallbacks(tools)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .defaultOptions(OpenAiChatOptions.builder()
                .internalToolExecutionEnabled(true).build())
                .build();

        ReactAgent  reactAgent = ReactAgent.builder()
                .name("ReAct Agent Demo")
                .chatClient(chatClient)
                .resolver(resolver)
                .maxIterations(10)
                .build();

        GraphRepresentation graphRepresentation = reactAgent.getStateGraph()
                .getGraph(GraphRepresentation.Type.PLANTUML);

        System.out.println("\n\n");
        System.out.println(graphRepresentation.content());
        System.out.println("\n\n");

        CompiledGraph compiledGraph = reactAgent.getAndCompileGraph();

        Optional<OverAllState> aiResult = compiledGraph.invoke(Map.of("messages", new UserMessage(question)));
        List<Message> messages = (List<Message>) aiResult.get().value("messages").get();
        AssistantMessage assistantMessage = (AssistantMessage) messages.get(messages.size() - 1);

        String answer = assistantMessage.getText();

        Map result = new HashMap<>();
        result.put("answer", answer);

        return result;
    }

    /**
     * mcp chat
     * @param requestMap
     * @return
     * @throws GraphStateException
     */
    @ApiDoc(desc = "mcpChat")
    @ShenyuDubboClient("/mcpChat")
    @Override
    public Map<String, Object> mcpChat(Map<String, Object> requestMap) throws GraphStateException {
        String tenantId = LoginInfoContextHelper.getTenantId();
        String question = MapUtils.getString(requestMap, "question");
        
        // 将tenantId信息拼接到question中，让AI知道当前服务的租户，特别是在调用MCP工具时
        String enhancedQuestion = String.format("当前租户ID：%s。在执行任何CRM操作时请使用此租户ID。用户问题：%s", tenantId, question);
        requestMap.put("tenantId", tenantId);

        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultToolCallbacks(tools)
                .build();

        System.out.println("\n>>> QUESTION: " + enhancedQuestion);
        System.out.println("\n>>> ASSISTANT: " + chatClient.prompt(enhancedQuestion).call().content());

        Map result = new HashMap<>();
        result.put("answer", chatClient.prompt(enhancedQuestion).call().content());

        return result;
    }

    /**
     * workflow chat
     * @param requestMap
     * @return
     * @throws GraphStateException
     */
    @ApiDoc(desc = "workflowChat")
    @ShenyuDubboClient("/workflowChat")
    @Override
    public Map<String, Object> workflowChat(Map<String, Object> requestMap) throws GraphStateException, GraphRunnerException {

        String question = MapUtils.getString(requestMap, "question");

        ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(new SimpleLoggerAdvisor()).build();

        OverAllStateFactory stateFactory = () -> {
            OverAllState state = new OverAllState();
            state.registerKeyAndStrategy("input", new ReplaceStrategy());
            state.registerKeyAndStrategy("classifier_output", new ReplaceStrategy());
            state.registerKeyAndStrategy("solution", new ReplaceStrategy());
            return state;
        };

        QuestionClassifierNode feedbackClassifier = QuestionClassifierNode.builder()
                .chatClient(chatClient)
                .inputTextKey("input")
                .categories(List.of("positive feedback", "negative feedback"))
                .classificationInstructions(
                        List.of("Try to understand the user's feeling when he/she is giving the feedback."))
                .build();

        QuestionClassifierNode specificQuestionClassifier = QuestionClassifierNode.builder()
                .chatClient(chatClient)
                .inputTextKey("input")
                .categories(List.of("after-sale service", "transportation", "product quality", "others"))
                .classificationInstructions(List
                        .of("What kind of service or help the customer is trying to get from us? Classify the question based on your understanding."))
                .build();

        StateGraph stateGraph = new StateGraph("Consumer Service Workflow Demo", stateFactory)
                .addNode("feedback_classifier", node_async(feedbackClassifier))
                .addNode("specific_question_classifier", node_async(specificQuestionClassifier))
                .addNode("recorder", node_async(new RecordingNode()))

                .addEdge(START, "feedback_classifier")
                .addConditionalEdges("feedback_classifier",
                        edge_async(new RecordingNode.FeedbackQuestionDispatcher()),
                        Map.of("positive", "recorder", "negative", "specific_question_classifier"))
                .addConditionalEdges("specific_question_classifier",
                        edge_async(new RecordingNode.SpecificQuestionDispatcher()),
                        Map.of("after-sale", "recorder", "transportation", "recorder", "quality", "recorder", "others",
                                "recorder"))
                .addEdge("recorder", END);

        GraphRepresentation graphRepresentation = stateGraph.getGraph(GraphRepresentation.Type.PLANTUML,
                "workflow graph");

        System.out.println("\n\n");
        System.out.println(graphRepresentation.content());
        System.out.println("\n\n");


        CompiledGraph compiledGraph = stateGraph.compile();

        String answer = compiledGraph.invoke(Map.of("input", question)).get().value("solution").get().toString();

        Map result = new HashMap<>();
        result.put("answer", answer);

        return result;
    }

    @ApiDoc(desc = "memoryChat")
    @ShenyuDubboClient("/memoryChat")
    @Override
    public Map<String, Object> memoryChat(Map<String, Object> requestMap) throws GraphStateException {

        String question = MapUtils.getString(requestMap, "question");
        String chatId = MapUtils.getString(requestMap, "chatId");

        ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(
                MessageChatMemoryAdvisor.builder(messageWindowChatMemory)
                        .build()
        ).build();

        chatClient.prompt(question).advisors(
                a -> a
                        .param(CONVERSATION_ID, chatId)
        ).stream().content();

        return null;
    }

    @ApiDoc(desc = "ragChat")
    @ShenyuDubboClient("/ragChat")
    @Override
    public Map<String, Object> ragChat(Map<String, Object> requestMap) throws GraphStateException {

        String question = MapUtils.getString(requestMap, "question");

        SearchRequest searchRequest = SearchRequest.builder().
                topK(4)
                .similarityThresholdAll()
                .query(question)
//                .filterExpression(new FilterExpressionBuilder().eq(textField, message).build())
                .build();

        // Step3 - Retrieve and llm generate
        String promptTemplate = getPromptTemplate(systemResource);
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(new RetrievalRerankAdvisor(vectorStore, rerankModel, searchRequest, new SystemPromptTemplate(promptTemplate), 0.1))
                .build();

        chatClient.prompt("用一个对话大方式,根据你知道大内容,有好的大回答用户的问题")
                .user(question)
                .stream()
                .chatResponse()
                .subscribe(chatResponse -> {
                    String content = chatResponse.getResult().getOutput().getText();
                    System.out.println("流式回复片段: " + content);
                    // 在这里处理每个文本片段
                }, error -> {
                    System.err.println("发生错误: " + error);
                }, () -> {
                    System.out.println("流已完成。");
                });

        return null;
    }

    private String getPromptTemplate(Resource systemResource) {
        try {
            return systemResource.getContentAsString(StandardCharsets.UTF_8);
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

}
